'''
Created on Apr 16, 2013

@author: drelfi
'''
from proxy.proxy_client import CacheProxyRequest, CacheProxyClientFactory
from proxy.certificates import certMaker
from twisted.web.proxy import Proxy
from twisted.internet.protocol import Protocol, ClientFactory
from twisted.web.http import  HTTPFactory
from twisted.internet import reactor, ssl, defer
from twisted.python import log
import urlparse
import sys
from OpenSSL import SSL

from OpenSSL import SSL as OSSL

"""
Since the define values in windows from a native library cannot be seen, they are hardcoded here
"""
SSL_CB_HANDSHAKE_START = 0x10
SSL_CB_HANDSHAKE_DONE = 0x20


     
class TunnelSSLContext(ssl.ContextFactory):
    def __init__(self, cert_filename, key_filename):
        self._cert_filename = cert_filename
        self._key_filename = key_filename
        
    def getContext(self):
        ctx = OSSL.Context(OSSL.SSLv23_METHOD)
        ctx.use_certificate_file(self._cert_filename)
        ctx.use_privatekey_file(self._key_filename)
        return ctx

class TunnelContextClientFactory(ssl.ClientContextFactory):
    
    def __init__(self, request):
        self.request = request
    
    def getContext(self):
        ctx = SSL.Context(SSL.TLSv1_METHOD)
        
        def _got_err(fail):
            log.err(fail)
            self.request.setResponseCode(500) 
            self.request.finish()
 
        def _got_cert(result):
            self.request.startSecureResponse(result["cert"], result["key"])
        
        def infoCallback(conn, where, ret):
            # conn is a OpenSSL.SSL.Connection
            # where is a set of flags telling where in the handshake we are
            # See http://www.openssl.org/docs/ssl/SSL_CTX_set_info_callback.html
            if where & SSL_CB_HANDSHAKE_DONE:
                #log.msg("Handshake done")
                peerCert = conn.get_peer_certificate()
                # Create the certificate for replying back
                r = certMaker(peerCert, self.request.host ).addCallback(_got_cert)
                r.addErrback(_got_err)
                
        ctx.set_info_callback( infoCallback )
        return ctx
    
class TunnelProtocol (Protocol): 
        """ 
        When a user agent makes a CONNECT request to a TunnelProxy, this 
        protocol implements the proxy's client logic. 
     
        When the proxy connects to the target host, it responds to the user 
        agent's request with an HTTP 200.  After that, it relays the stream 
        from the target host back through the connection to the user agent. 
     
        BUG: Verify that the 200 response meets the RFCs (or the common 
        practice if it deviates from the specification). 
        """ 
        # BUG: Handle early disconnects and other edge cases. 

        tunnelClient = None

        def connectionMade(self): 
            # Assigning 
            self.request.remoteProtocol = self
            
        def dataReceived(self, data):
            self.tunnelClient.dataReceived(data)
            
class TunnelClientFactory (ClientFactory): 
    
    protocol = TunnelProtocol
     
    def __init__(self, request):
        self.request = request
     
    def buildProtocol(self, addr):
        proto = ClientFactory.buildProtocol(self, addr)
        proto.request = self.request
        return proto
    
    def clientConnectionFailed(self, connector, reason): 
        self.request.setResponseCode(501, 'Gateway error') 
        self.request.finish() 

    
class TunnelProxyRequest (CacheProxyRequest): 
        protocols = {
            'http': CacheProxyClientFactory,
            'https': TunnelClientFactory
            }
        ports = {
            'http': 80,
            'https': 443
            }

        context = None
        
        """ 
        A request processor which supports the TUNNEL method. 
        """ 
        def process(self): 
            log.msg("Received a: %s - %d" % (self.method, id(self)) )
            if self.method.upper() == 'CONNECT':
                '''
                Opening a connect channel
                '''
                try: 
                    host, portStr = self.uri.split(':', 1) 
                    port = int(portStr) 
                except ValueError: 
                    # Either the connect parameter is not HOST:PORT or PORT is 
                    # not an integer, in which case this request is invalid. 
                    self.setResponseCode(400) 
                    self.finish() 
                else:  
                    self.context = TunnelContextClientFactory(self)
                    self.host = host
                    class_ = self.protocols["https"]
                    clientFactory = class_(self)
                    self.reactor.connectSSL(host, port, clientFactory, self.context)
            else:
                '''
                Difference between GET http://host:port/rest from GET /rest
                '''
                if self.isSecure():
                    log.msg("Secure GET request")
                    parsed = urlparse.urlparse(self.uri)
                    headers = self.getAllHeaders().copy()
                    self.host = headers["host"]
                    protocol = "http"
                    rest = urlparse.urlunparse(('', '') + parsed[2:])
                    self.clientFactory = self._process_secure_connect(self.host, protocol, rest)
                else:
                    '''
                    If it is GET http://host:port/rest
                    '''
                    log.msg("Normal GET request")
                    CacheProxyRequest.process(self)                 

        def startSecureResponse(self, certfile, keyfile ):
            log.msg("Connecting secure channel to browser")
            sslContext = TunnelSSLContext(certfile, keyfile)

            self.transport.write("HTTP/1.1 200 Connection established\r\n\r\n")
            self.transport.startTLS(sslContext)

        def _process_secure_connect(self, host, protocol, rest):
            '''
            Adapted twisted.web.proxy removing the connect call in order
            to do a secure connect when it is necessary.
            Use transport objects from previous request (CONNECT)
            '''
            if not rest:
                rest = rest + '/'
            class_ = self.protocols[protocol]
            headers = self.getAllHeaders().copy()
            if 'host' not in headers:
                headers['host'] = host
            self.content.seek(0, 0)
            s = self.content.read()
            
            # TODO - Check that requests[0] is a CONNECT request
            self.transport = self.channel.requests[0].transport
            
            clientFactory = class_(self.method, rest, self.clientproto, headers, s, self)
            
            # Run the protocol using transport objects from previous request
            protocol = clientFactory.buildProtocol(None)
            tunnelProtocol = self.channel.requests[0].remoteProtocol
            tunnelProtocol.tunnelClient = protocol
            protocol.makeConnection(tunnelProtocol.transport)

     
class TunnelProxy (Proxy): 
        """ 
        This class implements a simple web proxy with CONNECT support. 
     
        It creates a tunnel which can get the content passing through 
        """ 
        requestFactory = TunnelProxyRequest 
     
class TunnelProxyFactory (HTTPFactory): 
        """ 
        Factory for an HTTP proxy. 
     
        @ivar restrictedToPort: Only CONNECT requests to this port number 
            are allowed.  This may be None, in which case any port is allowed. 
        @type restrictedToPort: L{int} or None 
        """ 
     
        protocol = TunnelProxy 
     
        def __init__(self, logPath=None, timeout=60 * 60 * 12): 
            """ 
            @param logPath: The same as for HTTPFactory. 
            @param timeout: The same as for HTTPFactory. 
     
            """ 
            HTTPFactory.__init__(self, logPath, timeout) 
     
     
if __name__ == '__main__':
    
    reactor.listenTCP(4443, TunnelProxyFactory())
    log.startLogging(sys.stdout) # Log to the console
    reactor.run()
